'''
分割训练集，验证集，测试集
'''

import os
import sys

BASE_DIR = '/media/string/My Passport/CULane/CULane/'
ANNOTATIONS_DIR = 'Annotations/'
JPEG_DIR = 'JPEGImages/'
ALL_IMG_LIST = 'all_img_list.txt'
TRAIN_IMG_LIST = 'train_img_list.txt'
TRAIN_ANNO_LIST = 'train_anno_list.txt'
EVAL_IMG_LIST = 'eval_img_list.txt'
EVAL_ANNO_LIST = 'eval_anno_list.txt'
TEST_IMG_LIST = 'test_img_list.txt'

TOTAL_IMG_COUNT = 0


def FileCheck() -> None:
    '''
    检查图片和标注的完整性
    '''
    global TOTAL_IMG_COUNT
    f_ = open(os.path.join(BASE_DIR, ALL_IMG_LIST), 'w+')
    jpeg_dir_ = os.path.join(BASE_DIR, JPEG_DIR)
    vaild_file_count_ = 0
    total_file_count_ = 0

    for file_path_ in os.listdir(jpeg_dir_):
        total_file_count_ += 1
        file_path_img_ = os.path.join(JPEG_DIR, file_path_)
        file_path_anno_ = os.path.join(BASE_DIR, ANNOTATIONS_DIR, file_path_.replace('jpg', 'png'))
        # print(file_path_anno_)
        if (not os.path.isfile(file_path_anno_)):
            continue
        f_.write(file_path_.replace('.jpg', '') + '\n')
        vaild_file_count_ += 1
    f_.close()
    TOTAL_IMG_COUNT = total_file_count_
    print('所有的文件数量为：{0}，有效的文件数量为：{1}'.format(total_file_count_, vaild_file_count_))

def DataSegment() -> None:
    '''
    将图片分割为测试集，训练集，验证集
    '''
    global TOTAL_IMG_COUNT
    sl_ = open(os.path.join(BASE_DIR, ALL_IMG_LIST), 'r+').read().split('\n')
    index_ = 0
    train_img_list_path_ = os.path.join(BASE_DIR, TRAIN_IMG_LIST)
    train_anno_list_path_ = os.path.join(BASE_DIR, TRAIN_ANNO_LIST)
    eval_img_list_path_ = os.path.join(BASE_DIR, EVAL_IMG_LIST)
    eval_anno_list_path_ = os.path.join(BASE_DIR, EVAL_ANNO_LIST)
    test_img_list_path_ = os.path.join(BASE_DIR, TEST_IMG_LIST)
    f_train_img_list_ = open(train_img_list_path_, 'w+')
    f_train_anno_list_ = open(train_anno_list_path_, 'w+')
    f_eval_img_list_ = open(eval_img_list_path_, 'w+')
    f_eval_anno_list_ = open(eval_anno_list_path_, 'w+')
    f_test_img_list_ = open(test_img_list_path_, 'w+')

    for s_ in sl_:
        img_file_path_ = os.path.join(JPEG_DIR, s_ + '.jpg')
        anno_file_path_ = os.path.join(ANNOTATIONS_DIR, s_ + '.png')
        if (index_ < 70):
            # 训练集
            f_train_img_list_.write(img_file_path_ +  ' ' + anno_file_path_ + '\n')
            # f_train_anno_list_.write(anno_file_path_ + '\n')
        elif (70 <= index_ < 90):
            f_eval_img_list_.write(img_file_path_ + ' ' + anno_file_path_ + '\n')
            f_eval_anno_list_.write(anno_file_path_ + '\n')
        else:
            f_test_img_list_.write(img_file_path_ + '\n')
        index_ += 1
        if (index_ > 99):
            index_ = 0
    f_train_img_list_.close()
    f_train_anno_list_.close()
    f_eval_img_list_.close()
    f_eval_anno_list_.close()
    f_test_img_list_.close()


if __name__ == '__main__':
    FileCheck()
    DataSegment()
